package org.genitus.forceps.hbase;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.shaded.org.apache.commons.lang.exception.ExceptionUtils;
import org.apache.hadoop.hbase.shaded.org.apache.curator.shaded.com.google.common.collect.Lists;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.io.MD5Hash;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * Simple HBase clinet, only support BufferedMutator, Get and
 * Scan on sinlge (table, family, qualifier) currently.
 * <p>
 *   <b>
 *     THREAD-SAFE implementation, so that the client can be created once,
 *     and share it with different threads.
 *   </b>
 * </p>
 *
 * @author gwjiang (gwjiang@iflytek.com), 2016/9/8.
 */
public class HBaseClient {
  /** logger. */
  private static final Logger LOG = LoggerFactory.getLogger(HBaseClient.class);

  /** hbase connection. */
  private final Connection connection;

  /** hbase mutator, batched and asynchronous puts, thread-safe  */
  private BufferedMutator mutator = null;

  /** hbase table, light-weight and are not thread-safe, so make it final */
  private final Table table;

  /** column family in bytes. */
  private final byte[] col;

  /** colume qualifier in bytes. */
  private final byte[] quf;



  /**
   * Private ctor.
   *
   * @param zkQuorum zookeeper quorum.
   * @param tableName table name.
   * @param colFamily column family.
   * @param colQualifier column qualifier.
   * @throws IOException If there is an IO error.
   */
  private HBaseClient(String zkQuorum, String tableName, String colFamily , String colQualifier)
      throws IOException {
    try {
      col = Bytes.toBytes(colFamily);
      quf = Bytes.toBytes(colQualifier);

      Configuration config = HBaseConfiguration.create();
      config.set("hbase.zookeeper.quorum", zkQuorum);

      connection = ConnectionFactory.createConnection(config);
      table = connection.getTable(TableName.valueOf(tableName));
    } catch (Exception e) {
      LOG.error("# --------- gwjiang: initHandler failed");
      LOG.error(ExceptionUtils.getFullStackTrace(e));
      throw new IOException("hbase handler init failed");
    }
  }

  /**
   * Private ctor.
   *
   * @param zkQuorum zookeeper quorum.
   * @param tableName table name.
   * @param colFamily column family.
   * @param colQualifier column qualifier.
   * @param buffSize write buff size.
   * @throws IOException If there is an IO error.
   */
  private HBaseClient(String zkQuorum, String tableName, String colFamily , String colQualifier, int buffSize)
      throws IOException {
    this(zkQuorum, tableName, colFamily, colQualifier);

    BufferedMutatorParams params = new BufferedMutatorParams(table.getName());
    // set write buff size
    params.writeBufferSize(buffSize);
    mutator = connection.getBufferedMutator(params);
  }

  /**
   * Get a new HBaseClient.Builder instance.
   *
   * @param quorum zookeeper quorum.
   * @param table table name.
   * @param family column family.
   * @param qualifier column qualifier.
   * @return HBaseClient.Builder instance.
   */
  public static Builder newBuilder(String quorum, String table,
      String family, String qualifier) {
    return new Builder(quorum, table, family, qualifier);
  }

  /**
   * Hbase client builder.
   */
  public static class Builder {
    /** zookeeper quorum. */
    private String zkQuorum;

    /** table name. */
    private String table;

    /** column family. */
    private String colFamily;

    /** column qualifier. */
    private String colQualifier;

    /** write buff size  */
    private int buffSize = 1024 * 1024 * 6;

    /** enable write to table. */
    private boolean enableWrite = false;

    /**
     * ctor.
     *
     * @param quorum zookeeper quorum.
     * @param table table name.
     * @param family column family.
     * @param qualifier column qualifier.
     */
    public Builder(String quorum, String table, String family, String qualifier) {
      this.zkQuorum = quorum;
      this.table = table;
      this.colFamily = family;
      this.colQualifier = qualifier;
    }

    /**
     * Set write buff size.
     *
     * @param buffSize write buff size.
     * @return self.
     */
    public Builder withBuffSize(int buffSize) {
      this.buffSize = buffSize;
      return this;
    }

    /**
     * Enable write.
     *
     * @return self.
     */
    public Builder enableWrite() {
      this.enableWrite = true;
      return this;
    }

    /**
     * Disable write.
     *
     * @return self.
     */
    public Builder disableWrite() {
      this.enableWrite = false;
      return this;
    }

    /**
     * build.
     *
     * @return new instance.
     */
    public HBaseClient build() throws IOException {
      if (enableWrite) {
        return new HBaseClient(zkQuorum, table, colFamily, colQualifier, buffSize);
      } else {
        return new HBaseClient(zkQuorum, table, colFamily, colQualifier);
      }
    }
  }

  /**
   * Get data by rowkey.
   *
   * @param rowkey rowkey.
   * @return log data in bytes.
   */
  public ByteBuffer getLog(String rowkey) {
    LOG.debug("rowkey = {}", rowkey);

    try {
      Get get = new Get(Bytes.toBytes(rowkey));
      get.addColumn(col, quf);
      byte[] ret = table.get(get).getValue(col, quf);

      if (ret != null) {
        return ByteBuffer.wrap(ret);
      } else {
        return ByteBuffer.wrap(new byte[] {});
      }
    } catch (IOException e) {
      LOG.error("# --------- gwjiang: getLog failed");
      LOG.error(ExceptionUtils.getFullStackTrace(e));
      return ByteBuffer.wrap(new byte[] {});
    }
  }

  /**
   * Scan logs by seesion id or model id.
   *
   * @param id seesion id or model id.
   * @return log data list.
   */
  public List<LogData> getLogs(String id) {
    String prefix = MD5Hash.digest(id).toString().substring(0, 8) + "#" + id;
    String startKey = prefix + "#";
    String endKey = prefix + "$";

    LOG.debug("startKey = {}", startKey);
    LOG.debug("endKey   = {}", startKey);

    try {
      Scan scan = new Scan();
      scan.setStartRow(Bytes.toBytes(startKey));
      scan.setStopRow(Bytes.toBytes(endKey));

      List<LogData> logs = Lists.newArrayList();
      for (Result result : table.getScanner(scan)) {
        if (result != null && !result.isEmpty()) {
          byte[] value = result.getValue(col, quf);
          if (value.length != 0) {
            logs.add(new LogData(Bytes.toString(result.getRow()), ByteBuffer.wrap(value)));
          }
        }
      }

      return logs;
    } catch (Exception e) {
      LOG.error("# --------- gwjiang: getLogs failed\n" + ExceptionUtils.getFullStackTrace(e));
      return null;
    }
  }

  /**
   * Get all cells by rowkey.
   *
   * @param rowkey eg. traceId.
   * @return all cells.
   */
  public List<byte[]> getCells(String rowkey) {
    Get get = new Get(rowkey.getBytes());
    try {
      List<byte[]> list = Lists.newArrayList();
      list.addAll(table.get(get).listCells()
          .stream().map(CellUtil::cloneValue).collect(Collectors.toList()));

      return list;
    } catch (Exception e) {
      LOG.error("# --------- gwjiang: rowkey " + rowkey
          + " getCells failed\n" + ExceptionUtils.getFullStackTrace(e));
      return null;
    }
  }

  /**
   * Get data by rowkey.
   * Log with specific sid, there is only one.
   *
   * @param sid session id.
   * @return log data in bytes.
   */
  public ByteBuffer getSingleLog(String sid) {
    String prefix = MD5Hash.digest(sid).toString().substring(0, 8) + "#";
    String rowkey = prefix + sid;
    LOG.debug("rowkey = {}", rowkey);

    try {
      Get get = new Get(Bytes.toBytes(rowkey));
      get.addColumn(col, quf);
      byte[] ret = table.get(get).getValue(col, quf);

      if (ret != null) {
        return ByteBuffer.wrap(ret);
      } else {
        return ByteBuffer.wrap(new byte[] {});
      }
    } catch (IOException e) {
      LOG.error("# --------- gwjiang: getLog failed");
      LOG.error(ExceptionUtils.getFullStackTrace(e));
      return ByteBuffer.wrap(new byte[] {});
    }
  }

  /**
   * Write log with sid to hbase.
   *
   * @param sid session id.
   * @param data data.
   * @throws IOException If there is an IO error.
   */
  public void write(String sid, byte[] data) throws IOException {
    String prefix = MD5Hash.digest(sid).toString().substring(0, 8) + "#";
    String suffix = "#" + MD5Hash.digest(UUID.randomUUID().toString()).toString();
    String rowkey = prefix + sid + suffix;

    Put put = new Put(Bytes.toBytes(rowkey));
    put.addColumn(col, quf, data);
    put.setDurability(Durability.SKIP_WAL);

    // The mutations will be buffered and sent over the wire as part of a batch.
    mutator.mutate(put);
  }

  /**
   * Write log with sid to hbase with specific col.
   *
   * @param rowkey rowkey.
   * @param quf colume qualifier.
   * @param data data.
   * @throws IOException If there is an IO error.
   */
  public void write(String rowkey, String quf, byte[] data) throws IOException {
    Put put = new Put(Bytes.toBytes(rowkey));
    put.addColumn(col, Bytes.toBytes(quf), data);
    put.setDurability(Durability.SKIP_WAL);

    // The mutations will be buffered and sent over the wire as part of a batch.
    mutator.mutate(put);
  }

  /**
   * Write log with sid to hbase.
   * Log with specific sid, there is only one, so needn't suffix.
   *
   * @param sid session id.
   * @param data data.
   * @throws IOException If there is an IO error.
   */
  public void writeSingle(String sid, byte[] data) throws IOException {
    String prefix = MD5Hash.digest(sid).toString().substring(0, 8) + "#";
    String rowkey = prefix + sid;

    Put put = new Put(Bytes.toBytes(rowkey));
    put.addColumn(col, quf, data);
    put.setDurability(Durability.SKIP_WAL);

    // The mutations will be buffered and sent over the wire as part of a batch.
    mutator.mutate(put);
  }

  /**
   * Flush (only use in write), executes all the buffered, a
   * synchronous Mutation operations and waits until they are done.
   *
   * @throws IOException If there is an IO error.
   */
  public void flush() throws IOException {
    if (mutator != null) {
      mutator.flush();
    }
  }

  /**
   * Close client.
   *
   * @throws IOException If there is an IO error.
   */
  public void close() throws IOException {
    if (mutator != null) {
      mutator.close();
    }
    if (connection != null) {
      connection.close();
    }
  }
}
